package com.changan;

import ai.onnxruntime.OrtException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;

@Slf4j
@RestController
@RequestMapping("/images")
public class ImageController {

    private final MilvusService milvusService; // 自己封装 Milvus 插入/查询逻辑
    
    @Value("${image.upload.path:uploads}")
    private String uploadPath;
    
    @Value("${image.upload.url-prefix:/images}")
    private String urlPrefix;

    public ImageController(MilvusService milvusService) throws OrtException {
        this.milvusService = milvusService;
    }

    @PostMapping("/upload")
    public String upload(@RequestParam("file") MultipartFile file) throws Exception {
        // 保存图片到本地（也可以是OSS、MinIO）
        String fileName = UUID.randomUUID() + "_" + file.getOriginalFilename();
        Path savePath = Paths.get(uploadPath, fileName);
        Files.createDirectories(Paths.get(uploadPath));
        file.transferTo(savePath.toFile());

        log.info("开始处理上传文件: {}", fileName);
        
        // 获取向量
        float[] vector = milvusService.extractFeature(savePath.toString());
        
        log.info("文件 {} 特征提取完成", fileName);

        // 存入 Milvus
        milvusService.insertToMilvus(vector, urlPrefix + "/" + fileName);
        
        log.info("文件 {} 已存入 Milvus", fileName);

        return "上传成功: " + fileName;
    }
    
    // 添加搜索相似图片的接口
    @PostMapping("/search")
    public ResponseEntity<List<String>> searchSimilarImages(@RequestParam("file") MultipartFile file) throws Exception {
        // 保存临时文件用于特征提取
        String fileName = UUID.randomUUID() + "_" + file.getOriginalFilename();
        Path tempDir = Paths.get(System.getProperty("java.io.tmpdir"));
        Path savePath = tempDir.resolve(fileName);
        file.transferTo(savePath.toFile());
        
        log.info("开始处理搜索文件: {}", fileName);
        
        try {
            // 提取特征向量
            float[] vector = milvusService.extractFeature(savePath.toString());
            
            log.info("文件 {} 特征提取完成，开始在 Milvus 中搜索相似图片", fileName);
            
            // 在 Milvus 中搜索相似图片
            List<String> similarImages = milvusService.searchInMilvus(vector);
            
            log.info("在 Milvus 中找到 {} 个相似图片", similarImages.size());
            
            return ResponseEntity.ok(similarImages);
        } finally {
            // 删除临时文件
            Files.deleteIfExists(savePath);
            log.info("临时文件 {} 已删除", fileName);
        }
    }
    
    // 提供静态资源访问
    @GetMapping("/{fileName:.+}")
    public ResponseEntity<byte[]> getImage(@PathVariable String fileName) throws Exception {
        Path path = Paths.get(uploadPath, fileName);
        if (Files.exists(path)) {
            byte[] bytes = Files.readAllBytes(path);
            return ResponseEntity.ok(bytes);
        } else {
            return ResponseEntity.notFound().build();
        }
    }
}